Skip to content

Support TaskFlow call syntax on stub tasks for the Lang SDK - #69757

Open
jason810496 wants to merge 6 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag
Open

Support TaskFlow call syntax on stub tasks for the Lang SDK#69757
jason810496 wants to merge 6 commits into
apache:mainfrom
jason810496:feature/lang-sdk/taskflow-stub-dag

Conversation

@jason810496

@jason810496 jason810496 commented Jul 11, 2026

Copy link
Copy Markdown
Member

Scope

This PR ships the Python side only of the contract: parse-time capture of the TaskFlow call into a serialized arg-binding spec, the wire model, and its delivery to SDK runtimes through the Execution API and supervisor schema.
Nothing in this PR binds arguments inside a task runtime — the Go snippet below illustrates the consumer and lives in the stacked follow-up.

Stacked on top of this PR:

Why

@task.stub tasks can only be declared argless today, so a Go task that needs an upstream's output has to hand-write GetXCom calls (with the upstream task_id hard-coded in Go, duplicating the wiring the Dag file already expresses). This PR ships the Python side of making the natural TaskFlow call work across the language boundary:

@task.stub(queue="golang")
def transform(country: str, extracted: dict): ...


with DAG(...):
    transform("uk", extract())  # extract() is a normal Python @task
// The runtime (stacked #70209) binds "uk" onto country and pulls
// extract's XCom into extracted.
func transform(ctx sdk.TIRunContext, log *slog.Logger, country string, extracted map[string]any) error

Supported TaskFlow syntax

Every form below parses, serializes, and round-trips through the execution API (provider capture matrix in test_stub.py; the Dag shown is #70209's taskflow_binding_dag example):

@task.stub(queue="golang")
def via_flat_args(
    name: str,
    count: int,
    ratio: float,
    enabled: bool,
    tags: list,
    config: dict,
    numbers: list,
    note: str | None = None,
): ...


@dag(dag_id="taskflow_binding_dag")
def taskflow_binding_dag():
    via_flat_args(
        "summary",
        3,
        2.5,
        True,  # positional scalar literals (str/int/float/bool)
        ["metrics", "hourly"],  # array literal
        config=make_config(),  # keyword arg: XCom from another @task.stub
        numbers=make_numbers(),  # XCom binding onto a typed array parameter
    )  # `note` unpassed: its None default is captured as from_default
    region = make_region()
    via_struct_no_tags(RegionCode=region, Threshold=0.75)  # one XCom fanned into several calls
    via_struct_arg_tag(region_code=region, threshold=0.75)  # literal + XCom mixed as kwargs
    via_struct_unmatched_arg(region_code=region)  # defaulted param left unpassed
  • Literals of every JSON shape, positional or keyword; the wire value_schema is the JSON-schema fragment pydantic generates from the parameter annotation (TypeAdapter(annotation).json_schema() with a GenerateJsonSchema subclass layering OpenAPI's int64/double numeric formats): str{"type": "string"}, int{"type": "integer", "format": "int64"}, dict[str, int]{"type": "object", "additionalProperties": {...}}, list[int]{"type": "array", "items": {...}}, Literal["a", "b"]{"type": "string", "enum": [...]}, datetime/date/time/timedelta→string with the standard date-time/date/time/duration formats, unions→standard anyOf (str | None{"anyOf": [{"type": "string"}, {"type": "null"}]}); annotations pydantic cannot schema (arbitrary classes, unresolvable names) and untyped/Any parameters omit value_schema entirely (decode-only binding).
  • XComArg return values from any upstream — a normal @task or another @task.stub — which also wires the dependency edge (transform("uk", extract()) implies extract >> transform).
  • Keyword arguments normalize to declaration order through signature binding, so the serialized spec is always positional.
  • Defaults left unpassed are captured with from_default: true, letting keyword-style consumers (the Go sdk.TaskInput struct mode) leave them unclaimed.
  • Argless calls behave exactly as before: no spec is serialized, and pre-existing signatures (**kwargs, context-key parameter names) keep parsing.

Rejected loudly at parse time (v1 scope): .expand()/.partial() on a stub, stubs called with arguments inside a mapped task group, map/zip/concat XComArgs, indexing an upstream by custom XCom key, *args/**kwargs or context-key parameter names (only when the call actually passes arguments), and non-JSON-serializable literals (including NaN/Infinity).

How

  • Parse time (providers/standard): _StubOperator binds the TaskFlow call to the stub's signature and serializes an ordered positional-arg spec with the Dag (_arg_bindings); the cross-version imports it needs are routed through the common.compat sdk seam.
  • Wire model: each spec entry is one variant of a kind-discriminated union — XComArgBinding (pull of an upstream's return-value XCom by task_id) or LiteralArgBinding (inline JSON value, from_default-flagged when captured from a signature default) — carrying the stub parameter's name and an annotation-derived value_schema. The fragment is deliberately free-form (dict[str, JsonValue], not a typed model) so every keyword pydantic generated survives the server → supervisor → runtime trip verbatim; consumers validate the keywords they understand and ignore the rest, per JSON-schema semantics. The exact fragment shape follows the pydantic version active at Dag-parse time.
  • Server: ti_run returns the spec as a new optional TIRunContext.arg_bindings field, resolved through the shared DBDagBag only for _StubOperator tasks and stripped for older clients by a new execution-API version 2026-10-30 (Airflow 3.4 target); the supervisor schema gets a mirror version 2026-10-30 with a downgrade path, and the generated ts-sdk models plus the Go SDK's SupervisorSchemaVersion pin follow the schema bump in this PR. The serialized-Dag schema.json documents the optional per-task _arg_bindings property (no SERIALIZER_VERSION bump: optional field, no serialization-logic change).
  • Anything outside the v1 contract fails loudly at parse time (the rejected list above).

Was generative AI tooling used to co-author this PR?

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread go-sdk/dags/go_examples.py Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread go-sdk/pkg/binding/binding.go Outdated
Comment thread ts-sdk/src/generated/supervisor.ts

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will wait until next review then address my own comments to avoid CI-rerun.

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
@jason810496
jason810496 marked this pull request as ready for review July 22, 2026 09:17
@jason810496
jason810496 requested a review from uranusjr July 22, 2026 09:17

@ashb ashb left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall I like the direction. Almost all of my comments I can be challenged on, don't just make the changes if you think the current way is better/more correct

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment on lines +439 to +444
arg_bindings: list[TaskArgBinding] | None = None
"""
Ordered positional-argument binding spec for stub (foreign-runtime) tasks.

``None`` for regular tasks and for stub tasks that declare no parameters.
"""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmmmmm, I wonder if this should not allow none, and make it an empty list in that case. I don't think it functionally makes a difference but... 🤔

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if introducing arg_bindings should result in a bump in the serialization version? @amoghrajesh WDYT?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will update the airflow-core/src/airflow/serialization/schema.json to show the new args_binding field (the native Dag TaskFlow will leverage the args_binding field as well). However, I don't think it's not necessary to bump the "serialization version".

Comment on lines +3441 to +3443
assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or (
round_tripped.task_dict["extract"]._arg_bindings is None
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not hasattr or is None feels a bit odd. Pick one to assert.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given how tightly integrated into the exec API this change is, I'm not sure putting it in standard provider is right -- my first thought is that this should live in/with the dag parsing code, not with the stub operator code?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will look into whether it would be better to change the dag parsing side or to keep it as is.

class _StubOperator(DecoratedOperator):
custom_operator_name: str = "@task.stub"

# Mapped stubs would need per-map-index arg specs, which the foreign runtime cannot

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"Mapped stubs would need per-map-index arg specs"

I don't think this is true -- it's the same function for each mapped index (by design) so each index would receive the same type of arguments.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will double check this part.

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks you Ash for the review. I will address the comments shortly.

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py Outdated
assert response.status_code == 200
assert response.json()["arg_bindings"] == [
{"name": "country", "kind": "literal", "data_type": "string", "value": "uk"},
{"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"},

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, all the "data_type" here only comes from the annotation at the Stub Operator level. So what exactly the upstream task return doesn't really matter IMO.

There comes up another case that I needs to resolve. The case you give, user might annotate the argument as dict | bool, but I haven't deal with the union type annotation yet.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will update the airflow-core/src/airflow/serialization/schema.json to show the new args_binding field (the native Dag TaskFlow will leverage the args_binding field as well). However, I don't think it's not necessary to bump the "serialization version".

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will look into whether it would be better to change the dag parsing side or to keep it as is.

class _StubOperator(DecoratedOperator):
custom_operator_name: str = "@task.stub"

# Mapped stubs would need per-map-index arg specs, which the foreign runtime cannot

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will double check this part.

Comment thread providers/standard/tests/unit/standard/decorators/test_stub.py
Comment thread providers/standard/pyproject.toml Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py Outdated
@jason810496
jason810496 requested a review from Copilot July 24, 2026 08:22

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @ashb,

Here are the key updates since your last review:

  • Replaced the data_type enum with pydantic-generated JSON-schema fragments, carried per argument as value_schema.
  • Kept the argument materialization in the execution API:
    • In the existing Python world, the worker re-parses the Dag file and deserializes the operator, so it gets the call arguments for free.
    • A lang-SDK runtime can't parse Python, so it has to receive materialized bindings.
    • Doing this at Dag-processing time alone isn't enough -- resolving per-map-index values requires joining the TaskInstance at task runtime so ti_run in the execution API is the right place (it's also where API version negotiation strips the field for older clients).
  • Deferred mapped-operator support to #70570 and #70571 to keep this one concise enough to review.

@github-actions

github-actions Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

uv.lock on main just moved via #71039 ("Record the relaxed types-paramiko bound in uv.lock"), commit 381339d and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch 2 times, most recently from 7dd3bf9 to 2bafed2 Compare July 31, 2026 15:02
@uranusjr

uranusjr commented Aug 3, 2026

Copy link
Copy Markdown
Member

I can’t comment on the implementation, but the proposed interface looks very reasonable to me. It seems to me the described programming interface has not been fully implemented in this PR, and I don’t have enough knowledge to say whether this is on the right track or not tbh. Maybe it would be a good idea to edit the PR description to ground what exactly should be expected here.

@jason810496 jason810496 changed the title Support TaskFlow call syntax on stub tasks for the Go SDK Support TaskFlow call syntax on stub tasks for the Lang SDK Aug 4, 2026
@jason810496
jason810496 requested a review from ashb August 4, 2026 12:58
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch from 2bafed2 to 1b7e1d6 Compare August 4, 2026 13:03
@jason810496 jason810496 added this to the Airflow 3.4.0 milestone Aug 5, 2026
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch from 1b7e1d6 to 73bd2f8 Compare August 5, 2026 02:30
The @task.stub TaskFlow support in providers-standard imports
KNOWN_CONTEXT_KEYS, PlainXComArg, MappedOperator and the decorator base
classes through the compat layer so the provider keeps working down to
Airflow 2.11. Those symbols first ship in common-compat 1.19.0 (1.18.0
was released from main in the meantime without them), so the version is
cut here for the standard provider's pin to resolve.
Stub tasks silently ignored TaskFlow call arguments, so a Dag author
could not hand literals or upstream XCom results to a lang-SDK runtime.
The decorator now binds the call to the stub's signature at parse time
and captures an ordered arg spec (literal values and direct upstream
XCom references, with pydantic-derived JSON value schemas) that
serializes with the Dag, while rejecting what cannot cross the language
boundary: custom XCom keys, aggregated mapped outputs, non-JSON
literals, and stubs with arguments inside mapped task groups. Mapped
(.expand()) stubs capture no spec and keep the legacy behavior until a
follow-up delivers per-map-index bindings.
TIRunContext gains an arg_bindings field so a lang-SDK runtime receives
the stub task's TaskFlow arg spec at startup. ti_run derives it from the
serialized Dag only for stub operators, so regular tasks never pay for
the lookup, and only for clients on the new API version -- gated on the
Cadwyn VersionChangeWithSideEffects.is_applied check rather than a date
comparison -- so stub Dags that predate arg bindings keep running
against older clients, for which the version migration strips the field.
StartupDetails in the supervisor wire schema carries the new
arg_bindings so foreign runtimes receive the spec at task startup, with
a version migration that strips it for runtimes pinned to the previous
schema. The Go and TS SDKs regenerate against the new schema version;
the Go arg-binding runtime itself lands in a stacked follow-up PR.
An XComArg buried in a list or dict literal fell through to the JSON
check, whose "pass it in its JSON form instead" advice is impossible to
follow for a task output. Detect nested references up front and point
the author at the working alternative: pass the upstream output as its
own argument.
When a PR cuts a new provider version while the previous version is
still being voted on, only the rcN tags exist on the apache remote -
the final tag is pushed after the vote passes. The changes-table walk
in _get_all_changes_for_package assumed every past version has a final
tag and crashed with git exit 128 in that window, breaking CI for any
PR that bumps a provider version during a release wave.
@jason810496
jason810496 force-pushed the feature/lang-sdk/taskflow-stub-dag branch from 73bd2f8 to ed70fb2 Compare August 5, 2026 02:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: In progress

Development

Successfully merging this pull request may close these issues.

5 participants